import os
import random

random.seed(1)

trainval_percent = 1
train_percent = 0.8
# 读取标签txt
jsonfilepath = 'E:\\Test2\\all_text\\'

total_json = os.listdir(jsonfilepath)
num = len(total_json)
list = range(num)  # range(0,num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
# 检测文件夹是否创建
if not os.path.exists('E:\Test2\MYdata\\txt'):
    os.mkdir('E:\Test2\MYdata\\txt')

# 创建这些文件并往里面写入
ftrainval = open('E:\Test2\MYdata\\txt\\trainval.txt', 'w')
ftrain = open('E:\Test2\MYdata\\txt\\train.txt', 'w')
fval = open('E:\Test2\MYdata\\txt\\val.txt', 'w')

for i in list:
    print(i + 1)
    name = total_json[i].split('.')[0] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
# else:
#     ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
print('Create_txt Done, txt file in {} browser'.format('E:\Test2\MYdata\\txt'))
